#coding:utf-8
"""
saveTool
基于SQL注入，XSS的测试工具
@author:自之梦 10061661
@time: 17/10/11 上午11:30

"""
import sys
import threading 
import time
from loadTarget import loadTargetFile
from Queue import Queue


q = Queue()
def loadPoc(pocname):
	#__import__动态导入一个库
	try:
		pl = __import__(pocname)
	except:
		print "[Error]Load Poc failed..."
		sys.exit(1)
	try:
		pl = getattr(pl, pocname)
	except:
		print "[Error]Class %s is not existeds..." % pocname
		sys.exit(1)

	return pl


def saveFile(data):
	#将含有注入点的 URL写入文件
	Fname = str(time.time()) + ".txt"
	with open(Fname, "wb") as f:
		f.write(data)
	print "write url in %s success" % str(Fname)


def scan_thread(poc):
	while 1:
		target = q.get()
		print "[*]Testing %s" % target
		ret = poc.scan(target)
		if ret == True:
			#如果为真保存到文件中
			saveFile(target)
			print "[*]%s has vulnerablities..." % target
		else:
			return False
			print "[*] has not sql"


if __name__ == '__main__':
	try:
		filename = sys.argv[1]
		pocname = sys.argv[2]

	except:
		#target_file 文件文件
		#poc_name 插件名称
		print "[Usage] ./wvsf.py [target_file] [poc_name]"
		sys.exit(1)

	test_set = loadTargetFile(filename)
	for i in test_set:
		q.put(i)
	poc_class = loadPoc(pocname)
	poc = poc_class()

	print "[*]Starting scan thread..."
	poc = poc_class()
	for i in range(10):
		t = threading.Thread(target=scan_thread, args=(poc,))
 		t.start()